import random
import os

IMG_EXTENSIONS = [
    '.jpg',
    '.JPG',
    '.jpeg',
    '.JPEG',
    '.png',
    '.PNG',
    '.ppm',
    '.PPM',
    '.bmp',
    '.BMP',
]


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)

def get_image_paths(folder_path):
    image_paths = []
    # 遍历文件夹及其子文件夹
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # 检查文件扩展名是否是常见的图片格式
            if is_image_file(file):
                image_paths.append(os.path.join(root, file))
    return image_paths

def split_list_and_save(input_list, proportion, file1_name="list1.txt", file2_name="list2.txt"):
    """
    按照指定比例随机划分列表，并将结果保存到两个文本文件中。

    参数:
        input_list (list): 需要划分的原始列表。
        proportion (float): 第一个列表应占的比例（范围：0到1）。
        file1_name (str): 第一个列表保存的文件名（默认为list1.txt）。
        file2_name (str): 第二个列表保存的文件名（默认为list2.txt）。
    """
    if not (0 <= proportion <= 1):
        raise ValueError("proportion必须在0到1之间")

    # 打乱列表顺序
    random.shuffle(input_list)

    # 计算划分点
    split_index = int(len(input_list) * proportion)

    # 划分列表
    list1 = input_list[:split_index]
    list2 = input_list[split_index:]

    # 保存到文件
    with open(file1_name, "w", encoding="utf-8") as f1:
        for item in list1:
            f1.write(str(item) + "\n")
    with open(file2_name, "w", encoding="utf-8") as f2:
        for item in list2:
            f2.write(str(item) + "\n")

    print(f"第一个列表已保存到 {file1_name}，包含 {len(list1)} 个元素")
    print(f"第二个列表已保存到 {file2_name}，包含 {len(list2)} 个元素")

# 示例用法
if __name__ == "__main__":
    image_paths = get_image_paths("/hdd/sdg/wcj/datasets/OfficeHome")
    split_list_and_save(image_paths, proportion=0.8, file1_name="train_set.txt", file2_name="val_set.txt")